# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Common operations
"""

from inspect import signature
import numpy as np
import torch
from torch import nn, Tensor
import torch.nn.functional as F

from . import functions as func
from .functions import get_integer
from .units import Units, GLOBAL_UNITS

__all__ = [
    'GetVector',
    'GetDistance',
    'VelocityGenerator',
    'GetDistanceShift',
    'GetShiftGrad',
]


class GetVector(nn.Module):
    """The class to get vector with or without PBC box

    Args:
        use_pbc (bool): Whether to calculate vector under periodic boundary condition.
                       If `None` is given, it will determine whether to use periodic boundary
                       conditions based on whether the `pbc_box` is provided.
                       Default: None
    """
    def __init__(self, use_pbc: bool = None):
        super().__init__()
        self._use_pbc = use_pbc
        if use_pbc is None:
            self.calc_vector = self.calc_vector_default
        else:
            if use_pbc:
                self.calc_vector = self.calc_vector_pbc
            else:
                self.calc_vector = self.calc_vector_nopbc

    @property
    def use_pbc(self) -> bool:
        """whether to use periodic boundary condition"""
        return self._use_pbc

    @use_pbc.setter
    def use_pbc(self, use_pbc_: bool):
        """set whether to use periodic boundary condition"""
        self.set_pbc(use_pbc_)

    def set_pbc(self, use_pbc: bool):
        """set whether to use periodic boundary condition"""
        self._use_pbc = use_pbc
        if use_pbc is None:
            self.calc_vector = self.calc_vector_default
        else:
            if not isinstance(use_pbc, bool):
                raise TypeError(f'The type of use_pbc must be bool or None but got: {type(use_pbc)}')
            if use_pbc:
                self.calc_vector = self.calc_vector_pbc
            else:
                self.calc_vector = self.calc_vector_nopbc
        return self

    def calc_vector_default(self, initial: Tensor, terminal: Tensor, pbc_box: Tensor = None) -> Tensor:
        """get vector"""
        return func.calc_vector(initial, terminal, pbc_box)

    def calc_vector_pbc(self, initial: Tensor, terminal: Tensor, pbc_box: Tensor = None) -> Tensor:
        """get vector with perodic bundary condition"""
        return func.calc_vector_pbc(initial, terminal, pbc_box)

    def calc_vector_nopbc(self, initial: Tensor, terminal: Tensor, pbc_box: Tensor = None) -> Tensor:
        """get vector without perodic bundary condition"""
        return terminal - initial

    def forward(self, initial: Tensor, terminal: Tensor, pbc_box: Tensor = None):
        """Compute vector from initial point to terminal point."""
        return self.calc_vector(initial, terminal, pbc_box)


class GetDistance(GetVector):
    """Calculate distance with or without PBC box"""
    def __init__(self, use_pbc: bool = None, keepdims: bool = False, axis: int = -1):
        super().__init__(use_pbc=use_pbc)
        self.axis = get_integer(axis)
        self.keepdims = keepdims

    def forward(self, initial: Tensor, terminal: Tensor, pbc_box: Tensor = None):
        """Compute the distance from initial point to terminal point."""
        vector = self.calc_vector(initial, terminal, pbc_box)
        return torch.norm(vector, dim=self.axis, keepdim=self.keepdims)


class VelocityGenerator(nn.Module):
    """Generate velocities for atoms in system according to temperature"""
    def __init__(self,
                 temperature: float = 300,
                 remove_translation: bool = True,
                 seed: int = 0,
                 length_unit: str = None,
                 energy_unit: str = None,
                 ):
        super().__init__()

        if length_unit is None and energy_unit is None:
            self.units = GLOBAL_UNITS
        else:
            self.units = Units(length_unit, energy_unit)

        self.temperature = torch.tensor(temperature, dtype=torch.float32).reshape(-1, 1, 1)
        self.kb = torch.tensor(self.units.boltzmann, dtype=torch.float32)
        self.kbT = self.kb * self.temperature
        self.sigma = torch.sqrt(self.kbT)

        if seed is not None:
            torch.manual_seed(seed)

        self.kinectics_unit_scale = self.units.kinetic_ref
        self.remove_translation = remove_translation
        self.multi_temp = False

    def set_temperature(self, temperature: float):
        """set temperature"""
        self.temperature = torch.tensor(temperature, dtype=torch.float32).reshape(-1, 1, 1)
        self.multi_temp = len(self.temperature) > 1
        return self

    def forward(self, shape: tuple, atom_mass: Tensor, mask: Tensor = None):
        """Randomly generate velocities for atoms in system."""
        atom_mass = atom_mass.unsqueeze(-1)
        inv_mass = 1.0 / atom_mass
        velocity_scale = self.sigma * torch.sqrt(inv_mass / self.kinectics_unit_scale)
        
        if mask is not None:
            velocity_scale = torch.where(mask.unsqueeze(-1), velocity_scale, torch.zeros_like(velocity_scale))

        velocity = torch.randn(shape, device=atom_mass.device) * velocity_scale

        if self.remove_translation:
            momentum = atom_mass * velocity
            if mask is not None:
                dp = momentum.sum(dim=-2, keepdim=True) / mask.sum(dim=-1, keepdim=True).unsqueeze(-1)
            else:
                dp = momentum.mean(dim=-2, keepdim=True)
            momentum = momentum - dp
            velocity = momentum * inv_mass

        return velocity


class GetDistanceShift(nn.Module):
    """Calculate distance shifts"""
    def __init__(self, bonds: Tensor, num_atoms: int, num_walkers: int = 1, use_pbc: bool = None):
        super().__init__()
        self.bonds = bonds
        shape = (num_walkers, bonds.shape[-2], num_atoms)

        bond0 = self.bonds[..., 0].reshape(1, -1, 1).cpu().numpy()
        mask0 = np.zeros(shape)
        np.put_along_axis(mask0, bond0, 1, axis=-1)
        self.mask0 = torch.tensor(mask0, dtype=torch.int32).unsqueeze(-1)

        bond1 = self.bonds[..., 1].reshape(1, -1, 1).cpu().numpy()
        mask1 = np.zeros(shape)
        np.put_along_axis(mask1, bond1, 1, axis=-1)
        self.mask1 = torch.tensor(mask1, dtype=torch.int32).unsqueeze(-1)

        self.get_distance = GetDistance(use_pbc)

    def forward(self, coordinate_new: Tensor, coordinate_old: Tensor, pbc_box: Tensor = None):
        """Calculate distance shifts between new and old coordinates"""
        pos_new_0 = (self.mask0 * coordinate_new).sum(dim=-2)
        pos_new_1 = (self.mask1 * coordinate_new).sum(dim=-2)
        dis_new = self.get_distance(pos_new_0, pos_new_1, pbc_box)

        pos_old_0 = (self.mask0 * coordinate_old).sum(dim=-2)
        pos_old_1 = (self.mask1 * coordinate_old).sum(dim=-2)
        dis_old = self.get_distance(pos_old_0, pos_old_1, pbc_box)

        return dis_new - dis_old


class GetShiftGrad(nn.Module):
    """Calculate gradients of distance shifts"""
    def __init__(self, num_atoms: int, bonds: Tensor, num_walkers: int = 1, 
                 dimension: int = 3, use_pbc: bool = None):
        super().__init__()
        self.shape = (num_walkers, bonds.shape[-2], num_atoms, dimension)
        self.net = GetDistanceShift(bonds=bonds, num_atoms=num_atoms,
                                  num_walkers=num_walkers, use_pbc=use_pbc)
        self.zero_shift = torch.zeros((num_walkers, num_atoms - 1, num_atoms, dimension))

    def forward(self, coordinate_new: Tensor, coordinate_old: Tensor, pbc_box: Tensor = None):
        """Calculate gradients of distance shifts"""
        coordinate_new = coordinate_new.unsqueeze(1).expand(*self.shape)
        coordinate_old = coordinate_old.unsqueeze(1).expand(*self.shape)
        
        coordinate_new.requires_grad_(True)
        shift = self.net(coordinate_new, coordinate_old, pbc_box)
        shift_grad = torch.autograd.grad(shift.sum(), coordinate_new)[0]
        
        if torch.isnan(shift_grad).any():
            shift_grad = self.zero_shift.to(shift_grad.device)
            
        return shift_grad
